# src/data_access/data_handler.py (最终生产级版本)

from datetime import datetime, timedelta
from typing import List, Optional

import pandas as pd

from src.enums.timeframe import Timeframe
from .base_repository import BaseRepository
from ..core.indicators.indicator import Indicator


def get_max_indicator_period(indicators: List[Indicator]) -> int:
    """辅助函数：从指标请求列表中找出所需的最大计算周期。"""
    max_period = 0
    for indicator in indicators:
        # 假设指标的参数中有一个'period'或'length'键
        params = indicator.params
        period = params.get('period', params.get('length', 0))
        if period > max_period:
            max_period = period
    return max_period


class DataHandler:
    """
    一个通用的、由请求驱动的数据准备引擎 (生产级版本)。
    - 自动处理指标预热期
    - 健壮的多周期合并
    """

    def __init__(self, repository: BaseRepository):
        self.repo = repository
        self.master_df: pd.DataFrame = pd.DataFrame()
        self.base_timeframe: Optional[Timeframe] = None

    def prepare_data(
            self,
            symbol: str,
            timeframes: List[Timeframe],
            indicators: List[Indicator],
            start_date_str: str,
            end_date_str: str
    ) -> pd.DataFrame:

        print("--- [DataHandler] Starting data preparation ---")

        # 1. 【改进】计算预热期
        max_period = get_max_indicator_period(indicators)
        print(f"   Max indicator period detected: {max_period} bars.")

        # 假设基础周期是最小周期
        self.base_timeframe = min(timeframes, key=lambda tf: tf.to_milliseconds())
        # 估算预热期需要多长的时间 (乘以一个安全系数，例如1.5)
        # 这部分逻辑可以根据需要做得更精确
        warmup_delta = timedelta(milliseconds=self.base_timeframe.to_milliseconds() * max_period * 1.5)

        # 计算包含预热期的数据加载开始时间
        start_date_dt = datetime.fromisoformat(start_date_str)
        load_start_date = start_date_dt - warmup_delta

        print(f"   Loading data from {load_start_date} to include warm-up period.")

        # 2. 加载并安全地合并多周期数据 (使用新的加载开始时间)
        self._load_and_merge_timeframes(symbol, timeframes, load_start_date, datetime.fromisoformat(end_date_str))

        # 3. 动态计算所有需要的指标
        indicator_cols = self._calculate_indicators(indicators)

        # 4. 【改进】数据清理
        print("--- [DataHandler] Finalizing data... ---")

        # a. 只对新计算出的指标列执行dropna，或者如果您确信合并已经完美，
        #    可以对整个DataFrame执行。但为了安全，我们只对指标列操作。
        if indicator_cols:
            initial_rows = len(self.master_df)
            self.master_df.dropna(subset=indicator_cols, inplace=True)
            final_rows = len(self.master_df)
            print(f"   Dropped {initial_rows - final_rows} rows due to indicator warm-up NaNs.")

        # b. 【重要】在所有计算和清理完成后，再将DataFrame切片到用户请求的最终回测时间范围
        self.master_df = self.master_df.loc[start_date_str:end_date_str]

        print(f"   Final data sliced to requested range: {start_date_str} to {end_date_str}.")
        print(f"   Returning DataFrame with {len(self.master_df)} rows.")

        if self.master_df.empty:
            print(
                "WARNING: Final DataFrame is empty! Check data availability for the requested range and indicator "
                "periods.")

        print("--- [DataHandler] Data preparation complete. ---")
        return self.master_df

        # in src/data_access/data_handler.py

    def _load_and_merge_timeframes(
            self,
            symbol: str,
            timeframes: List[Timeframe],
            start_date: datetime,
            end_date: datetime
    ):
        if not timeframes:
            raise ValueError("必须至少提供一个时间周期。")

        print("1. Loading raw k-line data from repository...")
        self.base_timeframe = min(timeframes, key=lambda tf: tf.to_milliseconds())
        print(f"   Base timeframe set to: {self.base_timeframe.value}")

        # 加载所有数据，但这次我们不加前缀
        raw_data = {
            tf: self.repo.get_klines(symbol, tf, start_date, end_date)
            for tf in timeframes
        }

        # 【核心升级】为每个DataFrame创建多级索引列
        # 例如: ('1h', 'open'), ('1h', 'close')
        multi_index_dfs = {}
        for tf, df in raw_data.items():
            df.columns = pd.MultiIndex.from_product([[tf.value], df.columns])
            multi_index_dfs[tf] = df

        # 将基础周期的DataFrame作为主DataFrame
        self.master_df = multi_index_dfs[self.base_timeframe]
        self.master_df.sort_index(inplace=True)

        # 迭代并安全地合并其他周期
        for tf in sorted(timeframes, key=lambda t: t.to_milliseconds()):
            if tf == self.base_timeframe:
                continue

            print(f"   Securely merging timeframe: {tf.value}...")
            df_to_merge = multi_index_dfs[tf]
            df_to_merge.sort_index(inplace=True)

            # 【核心安全机制保持不变】
            # 先对大周期数据进行 .shift(1) 以模拟信息延迟
            df_shifted = df_to_merge.shift(1)

            # 使用 merge_asof 进行安全合并
            self.master_df = pd.merge_asof(
                left=self.master_df,
                right=df_shifted,
                left_index=True,
                right_index=True,
                direction='backward'
            )

        print("   Secure multi-timeframe merge complete.")

    def _calculate_indicators(self, indicators: List[Indicator]):
        new_columns = []
        # 这个方法的内部逻辑也保持不变
        # ... (与之前版本相同) ...
        # 2. Calculating indicators dynamically based on requests...
        print("2. Calculating indicators dynamically based on requests...")
        # 简单地按顺序执行，这要求用户在Factor中保证请求的顺序是正确的
        # (例如，先请求中间指标，再请求依赖于中间指标的最终指标)
        for indicator in indicators:
            output_col = indicator.column_name
            # 检查指标是否已经计算过（因为多个请求可能指向同一个指标）
            if output_col in self.master_df.columns:
                continue
            print(f"  - Calculating: {output_col}")
            # 【核心逻辑】调用指标对象自己的compute方法
            # 这种设计将DataHandler与任何具体的指标计算逻辑完全解耦
            try:
                output_series = indicator.compute(df=self.master_df)
                new_columns.append(output_col)
                self.master_df[output_col] = output_series
            except Exception as e:
                print(f"  - ERROR calculating '{output_col}': {e}")
                # 可以选择在这里抛出异常中断程序，或者继续计算其他指标
                # 抛出异常通常是更好的选择，因为它表明策略配置或数据有问题
                raise
        return new_columns
